"""
训练需要的词向量
"""
from chinese_text_classification.utils import read_data
from gensim.models import Word2Vec

path = "C:/Users/Administrator/PycharmProjects/pytorch/chinese_text_classification/THUCNews/data"
train_file = "train.txt"
dev_file = "dev.txt"

train_data, train_label = read_data(path, train_file)
dev_data, dev_label = read_data(path, dev_file)

data = train_data + dev_data
'''
data: 训练词向量需要的语料（确保已经分好词）
size: 词向量的维度大小（一般100-300）
window: 滑动窗口大小（语料长度长可以大一点，长度短就小一点）
min_count: 忽略词频小于这个数的单词
hs: 0代表词袋模型，1代表跳字模型，一般1的效果更好
workers: 并行的数目（需要装Cpython）
'''
model = Word2Vec(data, size=128, window=5, min_count=5, negative=3, sample=0.001, hs=1, workers=4)

# model.save('C:/Users/Administrator/PycharmProjects/pytorch/chinese_text_classification/THUCNews/saved_dict/word2vec')
model.wv.save_word2vec_format(
    'C:/Users/Administrator/PycharmProjects/pytorch/chinese_text_classification/THUCNews/saved_dict/word2vec.txt',
    binary=False)
# model.wv.save_word2vec_format(
#     'C:/Users/Administrator/PycharmProjects/pytorch/chinese_text_classification/THUCNews/saved_dict/word2vec.bin.gz',
#     binary=True)
